from PIL import Image
import random
from torchvision.transforms import functional as F
from torchvision.transforms import transforms as T


class Resize:
    def __init__(self, size):
        self.size = size

    def __call__(self, image, boxes):#lable [box]
        new_boxes = []
        w, h = image.size
        image = image.resize(self.size, Image.BILINEAR)
        ratio_x = self.size[0] / w
        ratio_y = self.size[1] / h
        for box in boxes:
            new_box = [
                box[0] * ratio_x,
                box[1] * ratio_y,
                box[2] * ratio_x,
                box[3] * ratio_y,                
            ]
            new_boxes.append(new_box)
        return image, new_boxes


class RandomFlip:
    def __init__(self, flip_prob = 0.5):
        self.flip_prob = flip_prob

    def __call__(self, image, boxes):
        if random.random() < self.flip_prob:
            new_boxes = []
            w, h = image.size
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
            for box in boxes:
                new_box = [
                    w - box[0] - box[2],
                    box[1],
                    box[2],
                    box[3],
                ]
                new_boxes.append(new_box)
            return image, new_boxes
        else:
            return image, boxes

class ToTensor:
    def __call__(self, image, boxes):
        size = image.size
        new_boxes = []
        for box in boxes:
            new_box = [
                box[0] / size[0],
                box[1] / size[1],
                box[2] / size[0],
                box[3] / size[1],
            ]
            new_boxes.append(new_box)
        image = F.to_tensor(image)
        return image, new_boxes


class Normalize:
    def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        self.mean = mean
        self.std = std

    def __call__(self, image, boxes):
        transform = T.Normalize(self.mean, self.std)
        image = transform(image)
        return image, boxes
    

class Compose:
    def __init__(self, transforms=None):
        self.transforms = transforms

    def __call__(self, image, boxes):
        for t in self.transforms:
            image, boxes = t(image, boxes)

        return image, boxes

            

if __name__ ==  '__main__':
    pass